# data_sources/akshare_source.py
import akshare as ak
import pandas as pd
from .base_data_source import BaseDataSource

class AkShareDataSource(BaseDataSource):
    def __init__(self):
        super().__init__("akshare")

    def get_stock_list(self) -> pd.DataFrame:
        try:
            stock_info = ak.stock_info_a_code_name()
            stock_info.columns = ['symbol', 'name']

            # 添加市场信息
            stock_info['market'] = stock_info['symbol'].apply(
                lambda x: 'SH' if x.startswith('6') else 'SZ'
            )
            stock_info['country'] = 'CN'

            return stock_info
        except Exception as e:
            self.logger.error(f"获取股票列表失败: {e}")
            return pd.DataFrame()

    def get_daily_data(self, symbol: str, start_date: str, end_date: str) -> pd.DataFrame:
        try:
            # 转换日期格式
            start_date = start_date.replace('-', '')
            end_date = end_date.replace('-', '')

            df = ak.stock_zh_a_hist(
                symbol=symbol,
                period="daily",
                start_date=start_date,
                end_date=end_date,
                adjust="hfq"
            )

            if df.empty:
                return pd.DataFrame()

            # 重命名列
            df.columns = ['trade_date', 'open_price', 'close_price', 'high_price',
                          'low_price', 'volume', 'turnover', 'amplitude',
                          'change_rate', 'turnover_rate', 'change_amount']

            df['symbol'] = symbol
            df['trade_date'] = pd.to_datetime(df['trade_date']).dt.strftime('%Y-%m-%d')

            return df

        except Exception as e:
            self.logger.error(f"AkShare数据获取失败 {symbol}: {e}")
            return pd.DataFrame()